#coding=utf8
from __future__ import division
from torch import nn
import torch
import torch.utils.data as torchdata
from torchvision import datasets,transforms
import os,time
import pandas as pd
import numpy as np
import torch.optim as optim
from torch.autograd import Variable
from torch.optim import lr_scheduler
from utils.train import train,trainlog
from sklearn.preprocessing import LabelEncoder
from torch.nn import CrossEntropyLoss
from twdata.twdataset import TWdata
from models.SE_Resnet1d import se_resnet20
from models.Resnet_1d import resnet20_1d
from models.Resnext_1d import resnext20_type1
from models.densenet_1d import densenet119
from sklearn.model_selection import train_test_split
from models.xception_1d import xception
from models.LeNet import LeNet,MLPNet
import logging
from twdata.twaugment import Compose, AddNoise, RandomAmplitude, DownSample, FlowNormalize, \
                        AddAxis, CenterCrop, RandomShiftCrop,ShiftCrop

from utils.ensemble_tools import prepare_data_split,oversample_by_num, class_to_idx_map, \
                         pred_aug, pred_single
from sklearn.model_selection import KFold, StratifiedKFold



class TWAug(object):
    def __init__(self):
        self.augment = Compose([
            AddNoise(A=0.1),
            RandomAmplitude(l=0.9,h=1.1),
            RandomShiftCrop(),
            FlowNormalize(),
            AddAxis()
        ])

    def __call__(self, spct):
        return self.augment(spct)

class TWAugVal(object):
    def __init__(self):
        self.augment = Compose([
            CenterCrop(),
            FlowNormalize(),
            AddAxis()
        ])

    def __call__(self, spct):
        return self.augment(spct)


class TWAugTest(object):
    def __init__(self, start_point,end_point ):
        self.augment = Compose([
            ShiftCrop(start_point=start_point,end_point=end_point),
            # CenterCrop(),
            FlowNormalize(),
            AddAxis()
        ])
    def __call__(self, spct):
        return self.augment(spct)


os.environ["CUDA_VISIBLE_DEVICES"] = "2"
n_splits = 8
epoch_num = 50
bs = 64*3
usecuda = 1
save_inter = 5
k_start = 0
start_epoch = 0


rawdata_root = '/media/gserver/data/tianwen/rawdata'
train_data_root = os.path.join(rawdata_root, 'first_train_data')
test_data_root = os.path.join(rawdata_root, 'first_test_data')


resume = None


# save dir
save_dir = '/media/gserver/models/tianwen/bagging/SEresnet20_crop/'
if not os.path.exists(save_dir):
    os.makedirs(save_dir)
logfile = '/media/gserver/models/tianwen/bagging/SEresnet20_crop/trainlog.log'
trainlog(logfile)


# k-fold bagging
kf = StratifiedKFold(n_splits=n_splits, random_state=137, shuffle=True)
k = 0

train_pd = pd.read_csv(os.path.join(rawdata_root, 'first_train_index_20180131.csv'))
test_pd = pd.read_csv(os.path.join(rawdata_root, 'first_test_index_20180131.csv'))

train_pd['type'] = train_pd['type'].apply(lambda x: class_to_idx_map[x])
test_pd['type'] = 0

for train_index, val_index in kf.split(train_pd, y=train_pd['type']):

    if k < k_start:
        k+=1
        continue

    logging.info('==='*40)
    logging.info('bagging training fold-%d'%k)

    train_fold = train_pd.loc[train_index,:]
    val_fold = train_pd.loc[val_index, :]

    train_fold = oversample_by_num(train_fold)

    # saving dir
    save_dir_k = os.path.join(save_dir, 'fold%d'%k)
    if not os.path.exists(save_dir_k):
        os.makedirs(save_dir_k)

    logging.info( 'train_fold label value_counts' )
    logging.info( train_fold.type.value_counts() )
    logging.info('val_fold label value_counts')
    logging.info( val_fold.type.value_counts() )


    data_set = {}
    data_set['train'] = TWdata(index_pd=train_fold,
                               data_root=train_data_root,
                               classes=class_to_idx_map.keys(),
                               transform=TWAug(),
                               )

    data_set['val'] = TWdata(index_pd=val_fold,
                             data_root=train_data_root,
                             classes=class_to_idx_map.keys(),
                             transform=TWAugVal(),
                             )

    logging.info( 'train augment:')
    for item in  data_set['train'].transform.augment.transforms:
        logging.info( '  %s %s'%(item.__class__.__name__, item.__dict__) )

    logging.info( 'val augment:' )
    for item in  data_set['val'].transform.augment.transforms:
        logging.info( '  %s %s'%(item.__class__.__name__, item.__dict__) )


    data_loader = {}
    data_loader['train'] = torchdata.DataLoader(data_set['train'], 32 * 3, num_workers=4,
                                                shuffle=True, pin_memory=True)
    data_loader['val'] = torchdata.DataLoader(data_set['val'], batch_size=128 * 3, num_workers=4,
                                              shuffle=False, pin_memory=True)


    # model prepare
    model = se_resnet20(num_classes=4)

    model = torch.nn.DataParallel(model)
    if resume:
        logging.info('resuming finetune from %s' % resume)
        model.load_state_dict(torch.load(resume))
        resume = None

    model = model.cuda()

    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-5)
    criterion = CrossEntropyLoss()
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.5)

    best_acc, best_model_wts = train(model,
                                     epoch_num,
                                     bs,
                                     start_epoch,
                                     optimizer,
                                     criterion,
                                     exp_lr_scheduler,
                                     data_set,
                                     data_loader,
                                     usecuda,
                                     save_inter,
                                     save_dir_k,
                                     print_inter=400,
                                     val_inter=4000)

    start_epoch = 0


    # pred test
    model.load_state_dict(best_model_wts)
    test_pred_k, _ = pred_aug(model,
                                 test_pd,
                                 test_data_root,
                                 class_to_idx_map,
                                 TWAugTest,
                                 bs,
                                 usecuda=1,
                                  num_crop=8,
                                  crop_size=0.8
                              )

    save_path = os.path.join(save_dir, 'fold%d-test.npy'%k)
    np.save(save_path, test_pred_k)
    logging.info('saved fold%d-test to %s' %(k,save_path))

    k+=1